from filter import heart_filter
import numpy as np
from scipy import signal
import pandas as pd
import correlate

max_module_num = 6
min_module_num = 4
correlate_threshold =0.30

sample_rate = 200

min_heart = int(0.6*sample_rate)
max_heart =int(2.0*sample_rate)

corr_step =int(sample_rate/20)

#计算差分
def differ(data ):
    gn1 = []
    for i in range(1, len(data) - 1):
        gn1.append((data[i + 1] - data[i - 1]) / 2)

    gn2 = []
    for i in range(2, len(data) - 2):
        gn2.append((2 * data[i + 1] + data[i + 2] - 2 * data[i - 1] - data[i - 2]) / 8)

    Gn = []
    for i in range(0, len(gn2)):
        Gn.append(gn1[i + 1] * gn1[i + 1] + gn2[i] * gn2[i])

    return  Gn

#list是点的位置，返回邻近点之间的间隔
def get_peak_dis(list):
    peak_dis =[]
    for i in range(len(list) - 1):
        peak_dis.append(list[i + 1] - list[i])
    return peak_dis

#获取模板的位置
def get_module_pos(acf):
    pattern_pos = []
    pos = 1
    for i in acf[1:]:
        if len(pattern_pos) > max_module_num:  # 模板数量上限
            break
        # 获取自相关后的位置
        if i > correlate_threshold and i > acf[pos - 1] and i > acf[pos + 1] and pos < 2 * sample_rate * (
                len(pattern_pos) + 1):
            if len(pattern_pos) == 0:
                pattern_pos.append(pos)
            else:
                if pos - pattern_pos[-1] > 0.5 * sample_rate and pos - pattern_pos[-1] < 2.0 * sample_rate:
                    pattern_pos.append(pos)
        pos = pos + 1
    return pattern_pos


#根据模板位置获取模板数据以及最大值点的位置
def get_module_data(pattern_pos,bcg):
    max_val_pos = []
    for i in range(len(pattern_pos)):
        if i == 0:
            sub_data = bcg[0:pattern_pos[i]]
            max_pos = sub_data.index(max(sub_data))
            max_val_pos.append(max_pos)
        else:
            sub_data = bcg[pattern_pos[i - 1]:pattern_pos[i]]
            max_pos = sub_data.index(max(sub_data)) + pattern_pos[i - 1]
            max_val_pos.append(max_pos)
    # print(max_val_pos)

    # 获取最大值之间的数据作为模板
    patter_data = []
    for i in range(0, len(max_val_pos) - 1):
        patter_data.append(bcg[max_val_pos[i]:max_val_pos[i + 1] - 1])

    return patter_data,max_val_pos


#循环计算候选窗口与模板的相关性，确定峰值点
def calculate_corr(bcg,max_val_pos,pattern_data):
    window_size = len(bcg)
    iteration_num = 0
    while max_val_pos[-1] + max_heart < window_size:
        iteration_num += 1
        if iteration_num > window_size / sample_rate * 2:
            break
        rre_of_window_length = []
        #从可能的最小到最大候选窗口进行计算
        for window_length in range(min_heart, max_heart, corr_step):
            sub_data = bcg[max_val_pos[-1]:max_val_pos[-1] + window_length]
            sub_data = pd.Series(sub_data)
            rrelation = []
            for i in pattern_data:
                if len(i) == 0:
                    continue
                pattern_i = signal.resample(i, window_length)
                pattern_i = pd.Series(pattern_i)
                cor = sub_data.corr(pattern_i)
                # 进行相关性计算
                # cor = correlate.dtw_distance(i,sub_data)
                rrelation.append(cor)
            max_rre = max(rrelation)  # 选择相关性最大的
            rre_of_window_length.append(max_rre)

        #选择与某个模板相关性最大的候选窗口，从中选择最大值点作为峰值点
        window_index = rre_of_window_length.index(max(rre_of_window_length))
        max_pos = int(max_val_pos[-1] + min_heart + window_index * corr_step)
        refine_data = bcg[max_val_pos[-1] + int(sample_rate / 2):max_pos]  # 避免太靠近上一个点
        refine_pos = []
        refine_val = []
        #选择极值点中的最大值
        for i in range(1, len(refine_data) - 1):
            if refine_data[i] > refine_data[i - 1] and refine_data[i] > refine_data[i + 1]:
                refine_pos.append(i)
                refine_val.append(refine_data[i])

        if len(refine_val) == 0:
            pos = signal.find_peaks(bcg[max_val_pos[-1] + int(sample_rate / 2):], distance=sample_rate / 2)
            pos = list(pos[0])
            if len(pos) == 0:
                break
            max_index = pos[0]
            max_pos = int(max_index + int(sample_rate / 2) + max_val_pos[-1] + 1)
            max_val_pos.append(max_pos)
            continue

        max_index = refine_val.index(max(refine_val))

        max_pos = int(refine_pos[max_index] + int(sample_rate / 2) + max_val_pos[-1] + 1)
        max_val_pos.append(max_pos)

    return max_val_pos

#将差分的峰值点对应到原始数据的峰值点，即把斜率最大的点对应到值最大的点
def map_differ_to_data(max_val_pos,bcg_filt):
    true_pos = []
    for i in range(len(max_val_pos)):
        cur_pos = max_val_pos[i]
        if bcg_filt[cur_pos] > bcg_filt[cur_pos - 1] and bcg_filt[cur_pos] > bcg_filt[cur_pos + 1]:
            true_pos.append(cur_pos)

        elif bcg_filt[cur_pos] > bcg_filt[cur_pos - 1] and bcg_filt[cur_pos] < bcg_filt[cur_pos + 1]:
            pos_i = cur_pos
            while bcg_filt[pos_i] < bcg_filt[pos_i + 1] and bcg_filt[pos_i] > bcg_filt[pos_i - 1]:
                pos_i -= 1
            left = pos_i - 1
            while bcg_filt[left] < bcg_filt[left - 1] and bcg_filt[left] > bcg_filt[left + 1]:
                left -= 1
            right = pos_i + 1
            while bcg_filt[right] < bcg_filt[right + 1] and bcg_filt[right] > bcg_filt[right - 1]:
                right += 1
            index = 0
            if bcg_filt[left] > bcg_filt[right]:
                index = left
            elif bcg_filt[left] <= bcg_filt[right]:
                index = right
            if len(true_pos) == 0:
                true_pos.append(index)
            elif len(true_pos) > 0 and left > true_pos[-1] + sample_rate / 2:
                true_pos.append(index)

        elif bcg_filt[cur_pos] < bcg_filt[cur_pos - 1] and bcg_filt[cur_pos] > bcg_filt[cur_pos + 1]:
            pos_i = cur_pos
            while bcg_filt[pos_i] < bcg_filt[pos_i - 1] and bcg_filt[pos_i] > bcg_filt[pos_i + 1]:
                pos_i += 1
            left = pos_i - 1
            while bcg_filt[left] < bcg_filt[left - 1] and bcg_filt[left] > bcg_filt[left + 1]:
                left -= 1
            right = pos_i + 1
            while bcg_filt[right] < bcg_filt[right + 1] and bcg_filt[right] > bcg_filt[right - 1]:
                right += 1
            index = 0
            if bcg_filt[left] > bcg_filt[right]:
                index = left
            elif bcg_filt[left] <= bcg_filt[right]:
                index = right
            if len(true_pos) == 0:
                true_pos.append(index)
            elif len(true_pos) > 0 and left > true_pos[-1] + sample_rate / 2:
                true_pos.append(index)
    return true_pos

#去除幅值异常的峰值点，超过3倍标准差
def remove_abnormal_val(max_val_pos,bcg_filt):
    val = [bcg_filt[i] for i in max_val_pos]
    mean_val = np.mean(val)
    std_val = np.std(val)
    true_pos1 = []
    for i in max_val_pos:
        if bcg_filt[i] > mean_val + 3 * std_val or bcg_filt[i] < mean_val - 3 * std_val:
            continue
        else:
            true_pos1.append(i)
    return true_pos1

#在两间隔过大的峰值点间尝试寻找峰值点，超过0.7倍标准差
def append_point(max_val_pos,bcg_filt):
    peak_dis = get_peak_dis(max_val_pos)
    mean_dis = np.mean(peak_dis)
    std_dis = np.std(peak_dis)
    val = [bcg_filt[i] for i in max_val_pos]
    std_val = np.std(val)
    min_dis = mean_dis + 0.7 * std_dis
    true_pos = []
    true_pos.append(max_val_pos[0])
    flag = 0
    while flag == 0 and max(peak_dis) > min_dis:
        for i in range(len(peak_dis)):
            if peak_dis[i] > min_dis:
                sub_data = list(bcg_filt[max_val_pos[i] + int(sample_rate / 2):max_val_pos[i + 1] - 1])
                sub_pos = signal.find_peaks(sub_data)[0]
                if len(sub_pos) == 0:
                    continue
                sub_val = [sub_data[i] for i in sub_pos]
                sub_index = sub_val.index(max(sub_val))

                candidate_pos = max_val_pos[i] + sub_pos[sub_index] + int(sample_rate / 2)

                if sub_val[sub_index] > np.min(val) - 3.0 * std_val \
                        and candidate_pos - max_val_pos[i] > sample_rate / 2 and max_val_pos[
                    i + 1] - candidate_pos > sample_rate / 2:
                    true_pos.append(max_val_pos[i] + sub_pos[sub_index] + int(sample_rate / 2))
                    flag = 0
                else:
                    flag = 1
            true_pos.append(max_val_pos[i + 1])
        max_val_pos = true_pos
        peak_dis = get_peak_dis(max_val_pos)
        true_pos = []
        if len(max_val_pos) <=0:
            break
        true_pos.append(max_val_pos[0])
    return max_val_pos


def get_heart_rate(bcg,sample_rate,low_cutoff):
    bcg_filt = heart_filter(bcg,1,low_cutoff,sample_rate)

    mean_val = np.mean(bcg_filt)
    bcg = [i - mean_val for i in bcg_filt]
    bcg = differ(bcg)#差分

    acf = np.correlate(bcg, bcg, mode='full')#自相关

    acf = acf[len(bcg) - 1:]
    acf = acf / acf[0]

    pattern_pos = get_module_pos(acf)#获取模板的位置

    #模板数量过少拒绝该窗口
    if pattern_pos== [] or len(pattern_pos) <min_module_num:
        return 0,[],[],[],0

    print(pattern_pos)
    ignore_num = len(pattern_pos)

    pattern_data,max_val_pos=get_module_data(pattern_pos,bcg)#根据模板位置获取模板数据以及最大值点的位置

    max_val_pos =calculate_corr(bcg,max_val_pos,pattern_data)#循环计算候选窗口与模板的相关性，确定峰值点

    max_val_pos =max_val_pos[ignore_num:]#去除模板的峰值点

    max_val_pos =map_differ_to_data(max_val_pos,bcg_filt)#将差分的峰值点对应到原始数据的峰值点

    max_val_pos =remove_abnormal_val(max_val_pos,bcg_filt)#去除幅值异常的峰值点

    max_val_pos =append_point(max_val_pos,bcg_filt)#在两间隔过大的峰值点间尝试寻找峰值点

    #获取峰值点间的间隔
    val = [bcg_filt[i] for i in max_val_pos]
    peak_dis = []
    for i in range(len(max_val_pos) - 1):
        peak_dis.append(max_val_pos[i + 1] - max_val_pos[i])

    heart_rate = 60 / (np.mean(peak_dis) / sample_rate)
    print(max_val_pos)

    return heart_rate,bcg_filt,max_val_pos,val,ignore_num